import backtrader as bt
import pandas as pd
import datetime


class three_bars(bt.Indicator):

    lines = ('up', 'down') # 明确有哪些line

    def __init__(self):
        self.addminperiod(4) # 最小周期
        self.plotinfo.plotmaster = self.data # 图根据哪个数据绘制的

    def next(self):
        self.up[0] = max(max(self.data.close.get(ago=-1, size=3)),
                               max(self.data.open.get(ago=-1, size=3)))  # 前三日最高价格
        self.down[0] = min(min(self.data.close.get(ago=-1, size=3)),
                           min(self.data.open.get(ago=-1, size=3)))  # 前三日最低价格


class MyStrategy(bt.Strategy):

    def __init__(self):
        self.up_down = three_bars(self.data)
        self.buy_signal = bt.indicators.CrossOver(self.data.close, self.up_down.up)  # 上穿
        self.sell_signal = bt.indicators.CrossDown(self.data.close, self.up_down.down)
        self.buy_signal.plotinfo.plot = False
        self.sell_signal.plotinfo.plot = False

    def next(self):
        if not self.position and self.buy_signal[0] == 1:  # 当前空仓且买入信号为1
            self.order = self.buy()

        if self.getposition().size < 0 and self.buy_signal[0] == 1:  # 当前为空头， 且买入信号为1
            self.order = self.close()
            self.order = self.buy()

        if not self.position and self.sell_signal[0] == 1:
            self.order = self.sell()

        if self.getposition().size > 0 and self.sell_signal[0] == 1:
            self.order = self.close()
            self.order = self.sell()

if __name__ == '__main__':
    cerebro = bt.Cerebro()
    dataframe = pd.read_csv("bitmex_data.csv")
    dataframe['open_time'] = pd.to_datetime(dataframe['open_time'])
    dataframe.set_index('open_time', inplace=True)
    dataframe['openinterest'] = 0
    brf_ = bt.feeds.PandasData(
        dataname=dataframe,
        fromdate=datetime.datetime(2020,10,19,4,24,00),
        todate=datetime.datetime(2020,10,19,12,43,00)
    )

    cerebro.adddata(brf_)
    cerebro.addstrategy(MyStrategy)

    cerebro.run()
    cerebro.plot()
